package top.yudoge.dogeeditor.llm.config.models

import dev.langchain4j.model.chat.ChatLanguageModel
import dev.langchain4j.model.chat.StreamingChatLanguageModel

class ModelRegistry {

    private val streamingChatModels = mutableMapOf<String, StreamingChatLanguageModel>()
    private val chatModels = mutableMapOf<String, ChatLanguageModel>()

    fun registerStreamingChatModel(name: String, model: StreamingChatLanguageModel) {
        streamingChatModels[name] = model
    }

    fun registerChatModel(name: String, model: ChatLanguageModel) {
        chatModels[name] = model
    }

    fun hasStreamingChatModel(name: String): Boolean {
        return streamingChatModels.containsKey(name)
    }

    fun getStreamingChatModel(name: String): StreamingChatLanguageModel {
        return streamingChatModels[name]?: throw IllegalArgumentException("Model $name not found")
    }

    fun hasChatModel(name: String): Boolean {
        return chatModels.containsKey(name)
    }

    fun getChatModel(name: String): ChatLanguageModel {
        return chatModels[name]?: throw IllegalArgumentException("Model $name not found")
    }

    fun allStreamingChatModelNames(): List<String> = streamingChatModels.keys.toList()

    fun allChatModelNames(): List<String> = chatModels.keys.toList()

}